import pandas as pd
from statsmodels.tsa.seasonal import seasonal_decompose
import matplotlib.pyplot as plt

# ==================== 配置参数 ====================
INPUT_FILE = "ori_message.xlsx"
OUTPUT_FILE = "monthly_average_data.xlsx"

# ==================== 数据处理函数（复用原代码） ====================
def process_data():
    excel_file = pd.ExcelFile(INPUT_FILE)
    years = ['2016', '2017', '2018', '2019', '2020', '2021']
    all_monthly_data = []
    for year in years:
        df = excel_file.parse(year)
        df['年'] = df['年'].ffill().astype(int)
        df['月'] = df['月'].ffill().astype(int)
        grouped = df.groupby(['年', '月']).agg({
            '流量(m3/s)': 'mean',
            '含沙量(kg/m3) ': 'mean'
        }).reset_index()
        grouped['排沙量(kg/s)'] = grouped['流量(m3/s)'] * grouped['含沙量(kg/m3) ']
        grouped = grouped.round({
            '流量(m3/s)': 2,
            '含沙量(kg/m3) ': 4,
            '排沙量(kg/s)': 2
        })
        all_monthly_data.append(grouped)
    result = pd.concat(all_monthly_data, ignore_index=True)
    result = result.sort_values(by=['年', '月']).reset_index(drop=True)
    result = result.rename(columns={'含沙量(kg/m3) ': '含沙量(kg/m3)'})
    result.to_excel(OUTPUT_FILE, index=False)
    return result

# ==================== 排沙量季节分解与绘图函数（关键修改） ====================
def plot_sediment_discharge_decompose(data):
    # 1. 构造完整日期索引（确保包含2021年12月）
    data['日期'] = pd.to_datetime(data[['年', '月']].apply(lambda x: f"{x['年']}-{x['月']:02d}-01", axis=1))
    data = data.set_index('日期')
    ts = data['排沙量(kg/s)']  # 提取排沙量时间序列

    # 2. 季节分解（乘法模型，周期12）
    decomposition = seasonal_decompose(ts, model='multiplicative', period=12)
    trend = decomposition.trend
    seasonal = decomposition.seasonal
    residual = decomposition.resid

    # 3. 创建子图（4行1列，修正横轴范围+残差散点图）
    plt.figure(figsize=(10, 12))

    # 子图1：原始排沙量（折线图，完整显示2017-2021年）
    plt.subplot(4, 1, 1)
    plt.plot(ts)
    plt.title('Multiplicative Decompose')
    plt.ylabel('排沙量(kg/s)')
    plt.xlim(pd.Timestamp('2017-01-01'), pd.Timestamp('2021-12-31'))  # 修正横轴结束时间为12月31日

    # 子图2：趋势成分（折线图，完整显示）
    plt.subplot(4, 1, 2)
    plt.plot(trend)
    plt.ylabel('Trend')
    plt.xlim(pd.Timestamp('2017-01-01'), pd.Timestamp('2021-12-31'))  # 修正横轴范围

    # 子图3：季节成分（折线图）
    plt.subplot(4, 1, 3)
    plt.plot(seasonal)
    plt.ylabel('Seasonal')
    plt.xlim(pd.Timestamp('2017-01-01'), pd.Timestamp('2021-12-31'))  # 统一横轴范围

    # 子图4：残差成分（散点图，关键修改）
    plt.subplot(4, 1, 4)
    plt.scatter(
        x=residual.index,  # 横轴为时间索引（包含2021年所有月份）
        y=residual.values,  # 纵轴为残差值
        s=15,  # 散点大小（稍大以便观察）
        alpha=0.7,  # 透明度
        color='darkblue'  # 颜色区分其他子图
    )
    plt.ylabel('Residual')
    plt.xlabel('年份')
    plt.xlim(pd.Timestamp('2017-01-01'), pd.Timestamp('2021-12-31'))  # 确保显示到12月

    plt.tight_layout()
    plt.savefig('排沙量时间序列分解图.png', dpi=300)
    plt.show()

# ==================== 主程序 ====================
if __name__ == "__main__":
    processed_data = process_data()
    plot_sediment_discharge_decompose(processed_data)
    print("\n========== 每月平均流量、含沙量及排沙量 ==========\n")
    print(processed_data.to_string(index=False))